# -*- coding:utf-8 -*-
import json
import os
import sys

import requests

__author__ = 'gin.chen'
__path__ = os.path.abspath(os.path.dirname(sys.argv[0])) + '/batch_config.json'


class livyOperator(object):
    def __init__(self, host, port, kind):
        self.kind = kind
        self.headers = {'Content-Type': 'application/json', 'X-Requested-By': 'admin'}
        self.sessions_url = None
        self.statements_url = None
        self.batch_url = None
        self.url = 'http://' + host + ':' + str(port)
        self.data = {'kind': self.kind}

    def create_session(self):
        r = requests.post(self.url + '/sessions', data=json.dumps(self.data), headers=self.headers)
        self.sessions_url = self.url + r.headers['location']
        print(r.json())

    def create_batches(self, path=__path__):
        try:
            with open(path) as data:
                r = requests.post(self.url + '/batches', data=data, headers=self.headers)
                self.batch_url = self.url + r.headers['location']
                print(r.json())
        except Exception as e:
            raise e

    def get_session(self):
        sessions_url = self.url + '/sessions'
        r = requests.get(sessions_url)
        for i in r.json().get('sessions'):
            if i.get('state') == 'idle' and i.get('kind') == self.kind:
                self.sessions_url = sessions_url + '/' + str(i.get('id'))
                break
        if self.sessions_url is None:
            self.create_session()
            return 1
        else:
            return 0
        # print(self.sessions_url)

    def show_session(self):
        r = requests.get(self.sessions_url)
        print(r.json())
        return r.json()

    def show_batches(self):
        r = requests.get(self.batch_url)
        # print(r.json())
        return r.json()

    def show_batch_log(self):
        r = requests.get(self.batch_url + '/log')
        print(r.json())

    def show_batches_all(self):
        r = requests.get(self.url + '/batches')
        print(r.json())

    def check_session_idle(self):
        session_status = self.show_session()
        return session_status.get('state') == 'idle'

    def execute_session(self, data):
        r = requests.post(self.sessions_url + '/statements', data=json.dumps({'code': data}), headers=self.headers)
        self.statements_url = self.url + r.headers['location']
        print(r.json())

    def show_statements(self):
        r = requests.get(self.statements_url)
        print(r.json())
        return r.json()

    def check_statements_finish(self):
        statements_context = self.show_statements()
        return statements_context.get('output') is not None

    def delete_session(self):
        requests.delete(self.sessions_url, headers=self.headers)
